#need RBTreeBASE.py
from RBTreeBASE import node,BST
class RBTree(BST):
    def __init__(self):
        self.NIL=node(key='NIL')
        self.root=self.NIL
    def inorder(self):
        BST.inorder(self,NIL=self.NIL)
    def rb_search(self,k):
        return self.search(k,NIL=self.NIL)
    def left_rotate(self,x):
        y=x.right
        x.right=y.left
        x.right.p=x
        y.p=x.p
        if y.p is self.NIL:
            self.root=y
        else:
            if x is x.p.left:
                x.p.left=y
            else:
                x.p.right=y
        y.left=x
        y.left.p=y
    def right_rotate(self,x):
        y=x.left
        x.left=y.right
        x.left.p=x
        y.p=x.p
        if y.p is self.NIL:
            self.root=y
        else:
            if x is x.p.left:
                x.p.left=y
            else:
                x.p.right=y
        y.right=x
        y.right.p=y
    def rb_insert(self,k):
        y=self.NIL
        x=self.root
        while x is not self.NIL:
            y=x
            if k<y.key:
                x=x.left
            else:
                x=x.right
        z=node(key=k,p=y,left=self.NIL,right=self.NIL,color='red')
        if y is self.NIL:
            self.root=z
        else:
            if k<y.key:
                y.left=z
            else:
                y.right=z
        self.rb_insert_fixup(z)
    def rb_insert_fixup(self,z):
        while z.p.color=='red': #if z.p is root, its color is black
            if z.p is z.p.p.left:
                y=z.p.p.right
                if y.color=='red':
                    z.p.color='black'
                    y.color='black'
                    z.p.p.color='red'
                    z=z.p.p
                else:
                    if z is z.p.right:
                        z=z.p
                        self.left_rotate(z)
                    z.p.color='black'
                    z.p.p.color='red'
                    self.right_rotate(z.p.p)
            else: #z.p is z.p.p.right
                y=z.p.p.left
                if y.color=='red':
                    z.p.color='black'
                    y.color='black'
                    z.p.p.color='red'
                    z=z.p.p
                else:
                    if z is z.p.left:
                        z=z.p
                        self.right_rotate(z)
                    z.p.color='black'
                    z.p.p.color='red'
                    self.left_rotate(z.p.p)
        self.root.color='black'
    def rb_delete(self,z):
        if z.left is self.NIL or z.right is self.NIL:
            y=z
        else:
            y=z.successor(NIL=self.NIL)
        if y.left is not self.NIL:
            x=y.left
        else:
            x=y.right
        x.p=y.p
        if y.p is self.NIL:
            self.root=x
        else:
            if y is y.p.left:
                y.p.left=x
            else:
                y.p.right=x
        if y is not z:
            z.key=y.key
        if y.color=='black':
            self.rb_delete_fixup(x)
    def rb_delete_fixup(self,x):
        while x is not self.root and x.color=='black':
            if x is x.p.left:
                w=x.p.right
                if w.color=='red':
                    w.color='black'
                    x.p.color='red'
                    self.left_rotate(x.p)
                    w=x.p.right
                if w.left.color=='black' and w.right.color=='black':
                    w.color='red'
                    x=x.p
                elif w.right.color=='black':
                    w.left.color='black'
                    w.color='red'
                    self.right_rotate(w)
                    w=x.p.right
                w.color=x.p.color
                x.p.color='black'
                w.right.color='black'
                self.left_rotate(x.p)
                x=self.root
            else: #x is x.p.right
                w=x.p.left
                if w.color=='red':
                    w.color='black'
                    x.p.color='red'
                    self.right_rotate(x.p)
                    w=x.p.left
                if w.right.color=='black' and w.left.color=='black':
                    w.color='red'
                    x=x.p
                elif w.left.color=='black':
                    w.right.color='black'
                    w.color='red'
                    self.left_rotate(w)
                    w=x.p.left
                w.color=x.p.color
                x.p.color='black'
                w.left.color='black'
                self.right_rotate(x.p)
                x=self.root
        x.color='black'
    
if __name__=='__main__':
    tree=RBTree();
    A=[15,5,16,3,12,20,10,13,18,23,6,7]
    #A=[3,1,2]
    for i in A:
        tree.rb_insert(i)
    print 'tree inorder:'
    tree.inorder()
    z=tree.rb_search(12)
    if z is not tree.NIL:
        tree.rb_delete(z)
    print 'deleted 12:'
    tree.inorder()


